#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Neo4j知识图谱构建和自然语言查询工具
功能：将决策树数据存储到Neo4j，支持自然语言查询
"""

from neo4j import GraphDatabase
import jieba
import re
import json
from typing import List, Dict, Any


class Neo4jKnowledgeGraph:
    def __init__(self, uri="bolt://localhost:7687", user="neo4j", password="password"):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
        
        # 同义词映射
        self.synonyms = {
            '无产物': ['没有产物', '无条带', '扩增失败', '无扩增产物'],
            '产物量少': ['产物量很少', '条带弱', '扩增效率低', '产出低'],
            '模板': ['DNA模板', 'RNA模板', '起始模板', '反应模板'],
            '质量': ['品质', '完整性', '纯度'],
            '降解': ['破碎', '断裂', '不完整', '质量差'],
            '建议': ['推荐', '应该', '可以', '需要'],
            '鉴定': ['检测', '检查', '验证', '分析']
        }
    
    def close(self):
        """关闭数据库连接"""
        self.driver.close()
    
    def clear_database(self):
        """清空数据库"""
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
            print("✅ 数据库已清空")
    
    def parse_decision_path(self, text: str) -> Dict:
        """解析决策路径"""
        # 分割路径
        parts = text.split(' -> ')
        
        if len(parts) < 2:
            return None
        
        # 提取产品编号
        product = parts[0].strip()
        
        # 提取货号
        cargo_match = re.search(r'\[货号:([^\]]+)\]', text)
        cargo_number = cargo_match.group(1) if cargo_match else product
        
        # 构建决策路径
        path_nodes = []
        for i, part in enumerate(parts[1:], 1):
            # 清理货号信息
            clean_part = re.sub(r'\[货号:[^\]]+\]', '', part).strip()
            if clean_part:
                path_nodes.append({
                    'level': i,
                    'content': clean_part,
                    'type': self._classify_node_type(clean_part, i)
                })
        
        return {
            'product': product,
            'cargo_number': cargo_number,
            'path_nodes': path_nodes,
            'full_text': text
        }
    
    def _classify_node_type(self, content: str, level: int) -> str:
        """分类节点类型"""
        if '？' in content or '?' in content:
            return 'question'
        elif '建议' in content or '推荐' in content:
            return 'solution'
        elif '是' == content or '否' == content:
            return 'answer'
        elif level == 1:
            return 'problem'
        elif level == 2:
            return 'category'
        elif '则' in content:
            return 'condition'
        else:
            return 'step'
    
    def create_knowledge_graph(self, decision_paths: List[str]):
        """创建知识图谱"""
        with self.driver.session() as session:
            for path_text in decision_paths:
                parsed = self.parse_decision_path(path_text)
                if not parsed:
                    continue
                
                # 创建产品节点
                session.run("""
                    MERGE (p:Product {code: $product, cargo_number: $cargo_number})
                    SET p.name = $product
                """, product=parsed['product'], cargo_number=parsed['cargo_number'])
                
                # 创建路径节点并建立关系
                prev_node_id = f"product_{parsed['product']}"
                
                for i, node in enumerate(parsed['path_nodes']):
                    node_id = f"{parsed['product']}_level_{node['level']}_{hash(node['content'])}"
                    
                    # 创建节点
                    session.run(f"""
                        MERGE (n:Node {{
                            id: $node_id,
                            content: $content,
                            type: $type,
                            level: $level,
                            product: $product
                        }})
                        SET n.keywords = $keywords
                    """, 
                    node_id=node_id,
                    content=node['content'],
                    type=node['type'],
                    level=node['level'],
                    product=parsed['product'],
                    keywords=' '.join(jieba.cut(node['content']))
                    )
                    
                    # 建立关系
                    if i == 0:
                        # 产品到第一个节点
                        session.run("""
                            MATCH (p:Product {code: $product})
                            MATCH (n:Node {id: $node_id})
                            MERGE (p)-[:HAS_PROBLEM]->(n)
                        """, product=parsed['product'], node_id=node_id)
                    else:
                        # 节点之间的关系
                        session.run("""
                            MATCH (prev:Node {id: $prev_id})
                            MATCH (curr:Node {id: $curr_id})
                            MERGE (prev)-[:LEADS_TO]->(curr)
                        """, prev_id=prev_node_id, curr_id=node_id)
                    
                    prev_node_id = node_id
                
                print(f"✅ 已处理: {parsed['product']} - {len(parsed['path_nodes'])} 个节点")
    
    def natural_language_query(self, query: str) -> List[Dict]:
        """自然语言查询"""
        # 提取查询关键词
        keywords = self._extract_query_keywords(query)
        print(f"🔍 提取的关键词: {keywords}")
        
        results = []
        
        with self.driver.session() as session:
            # 策略1: 产品精确匹配
            if 'products' in keywords:
                for product in keywords['products']:
                    cypher = """
                        MATCH (p:Product {code: $product})-[:HAS_PROBLEM*1..10]->(n:Node)
                        RETURN p.code as product, n.content as content, n.type as type, n.level as level
                        ORDER BY n.level
                    """
                    result = session.run(cypher, product=product)
                    for record in result:
                        results.append({
                            'strategy': 'product_match',
                            'product': record['product'],
                            'content': record['content'],
                            'type': record['type'],
                            'level': record['level'],
                            'score': 10
                        })
            
            # 策略2: 内容关键词匹配
            if 'terms' in keywords:
                for term in keywords['terms'][:3]:  # 限制关键词数量
                    cypher = """
                        MATCH (n:Node)
                        WHERE n.content CONTAINS $term OR n.keywords CONTAINS $term
                        MATCH (p:Product)-[:HAS_PROBLEM*1..10]->(n)
                        RETURN p.code as product, n.content as content, n.type as type, n.level as level
                        ORDER BY n.level
                    """
                    result = session.run(cypher, term=term)
                    for record in result:
                        results.append({
                            'strategy': 'content_match',
                            'product': record['product'],
                            'content': record['content'],
                            'type': record['type'],
                            'level': record['level'],
                            'score': 7
                        })
            
            # 策略3: 问题类型匹配
            if 'terms' in keywords:
                for term in keywords['terms']:
                    cypher = """
                        MATCH (n:Node {type: 'problem'})
                        WHERE n.content CONTAINS $term
                        MATCH (p:Product)-[:HAS_PROBLEM]->(n)
                        MATCH path = (n)-[:LEADS_TO*0..5]->(solution:Node)
                        WHERE solution.type = 'solution'
                        RETURN p.code as product, n.content as problem, solution.content as solution
                    """
                    result = session.run(cypher, term=term)
                    for record in result:
                        results.append({
                            'strategy': 'problem_solution',
                            'product': record['product'],
                            'problem': record['problem'],
                            'solution': record['solution'],
                            'score': 9
                        })
            
            # 策略4: 解决方案匹配
            if 'terms' in keywords:
                for term in keywords['terms']:
                    cypher = """
                        MATCH (n:Node {type: 'solution'})
                        WHERE n.content CONTAINS $term
                        MATCH (p:Product)-[:HAS_PROBLEM*1..10]->(n)
                        RETURN p.code as product, n.content as solution
                    """
                    result = session.run(cypher, term=term)
                    for record in result:
                        results.append({
                            'strategy': 'solution_match',
                            'product': record['product'],
                            'solution': record['solution'],
                            'score': 8
                        })
        
        # 去重和排序
        unique_results = self._deduplicate_results(results)
        return sorted(unique_results, key=lambda x: x['score'], reverse=True)
    
    def _extract_query_keywords(self, query: str) -> Dict:
        """提取查询关键词"""
        # 产品编号
        products = re.findall(r'[A-Z]{1,3}\d{3,4}', query.upper())
        
        # 分词并过滤
        words = list(jieba.cut(query))
        terms = []
        
        for word in words:
            word = word.strip()
            if len(word) > 1 and word not in ['的', '是', '有', '和', '怎么', '什么']:
                terms.append(word)
                
                # 添加同义词
                for key, synonyms in self.synonyms.items():
                    if word in synonyms or word == key:
                        terms.extend([key] + synonyms)
        
        return {
            'products': list(set(products)),
            'terms': list(set(terms))
        }
    
    def _deduplicate_results(self, results: List[Dict]) -> List[Dict]:
        """去重结果"""
        seen = set()
        unique_results = []
        
        for result in results:
            # 创建唯一键
            if 'solution' in result:
                key = f"{result.get('product', '')}-{result.get('solution', '')}"
            else:
                key = f"{result.get('product', '')}-{result.get('content', '')}"
            
            if key not in seen:
                seen.add(key)
                unique_results.append(result)
        
        return unique_results
    
    def get_decision_path(self, product: str, problem: str) -> List[Dict]:
        """获取完整决策路径"""
        with self.driver.session() as session:
            cypher = """
                MATCH (p:Product {code: $product})-[:HAS_PROBLEM]->(start:Node)
                WHERE start.content CONTAINS $problem
                MATCH path = (start)-[:LEADS_TO*0..10]->(end:Node)
                RETURN nodes(path) as path_nodes
                ORDER BY length(path) DESC
                LIMIT 1
            """
            result = session.run(cypher, product=product, problem=problem)
            
            for record in result:
                path_nodes = []
                for node in record['path_nodes']:
                    path_nodes.append({
                        'content': node['content'],
                        'type': node['type'],
                        'level': node['level']
                    })
                return path_nodes
        
        return []
    
    def display_results(self, results: List[Dict], query: str):
        """显示查询结果"""
        if not results:
            print("\n❌ 没有找到相关结果")
            return
        
        print(f"\n🎯 查询结果 (共{len(results)}条):")
        print("=" * 80)
        
        for i, result in enumerate(results[:10], 1):  # 最多显示10条
            print(f"\n--- 结果 {i} ---")
            print(f"🏷️  产品: {result.get('product', 'N/A')}")
            print(f"🔍 匹配策略: {result.get('strategy', 'N/A')}")
            print(f"📊 相关度: {result.get('score', 0)}")
            
            if 'problem' in result and 'solution' in result:
                print(f"❓ 问题: {result['problem']}")
                print(f"💡 解决方案: {result['solution']}")
            elif 'solution' in result:
                print(f"💡 解决方案: {result['solution']}")
            elif 'content' in result:
                print(f"📝 内容: {result['content']} (类型: {result.get('type', 'N/A')})")


def demo_neo4j_knowledge_graph():
    """演示Neo4j知识图谱"""
    # 示例数据
    sample_data = [
        "PK511 -> 无产物或产物量很少 -> 模板 -> 是否鉴定过模板质量 -> 是 -> 质量较好，则没有问题 [货号:PK511]",
        "PK511 -> 无产物或产物量很少 -> 模板 -> 是否鉴定过模板质量 -> 是 -> 模板降解，则可能扩增不出条带 -> 建议重新制备模板 [货号:PK511]",
        "PK511 -> 无产物或产物量很少 -> 模板 -> 是否鉴定过模板质量 -> 否 -> 建议鉴定模板质量 -> cDNA不可直接跑胶，可以跑胶鉴定RNA的质量 [货号:PK511]",
        "PK511 -> 无产物或产物量很少 -> 模板 -> 是否鉴定过模板质量 -> 否 -> 建议鉴定模板质量 -> 质粒/gDNA可以跑胶鉴定，上样量200-400ng [货号:PK511]"
    ]
    
    # 创建知识图谱
    kg = Neo4jKnowledgeGraph()
    
    try:
        print("🗃️  清空数据库...")
        kg.clear_database()
        
        print("🏗️  构建知识图谱...")
        kg.create_knowledge_graph(sample_data)
        
        print("\n🔍 测试自然语言查询...")
        
        # 测试查询
        test_queries = [
            "PK511无产物怎么办",
            "模板质量问题",
            "建议重新制备",
            "cDNA检测方法"
        ]
        
        for query in test_queries:
            print(f"\n{'='*60}")
            print(f"查询: {query}")
            results = kg.natural_language_query(query)
            kg.display_results(results, query)
    
    finally:
        kg.close()


if __name__ == "__main__":
    demo_neo4j_knowledge_graph()
